#ifndef ML_DECISION_STUMP_H_
#define ML_DECISION_STUMP_H_

#include "ML_WeakLearner.h"
#include "ML_SubclassFolder.h"
#include <set>

GLOBAL_NAMESPACE_BEGIN
NAMESPACE_MACHINE_LEARNING_BEGIN

class MLDecisionStumpParams;
class MLDecisionStumpImpl;
class MLDecisionStumpPreprocess;

class MLDecisionStump : public MLWeakLearner
{
public:
    MLDecisionStump();
    MLDecisionStump(const MLDecisionStump& rhs);
    ~MLDecisionStump();

    void reset();
    //void setSubClasses(const std::set<int>& classes, int maxLabel, int thisClassIdx);
    void setSubClassFolder(MLSubClassFolder* subClasses);
    void setPreprocessData(MLDecisionStumpPreprocess* preProcess);

    int     train    (const Eigen::MatrixXd& trainData, const Eigen::MatrixXi& labels, const Eigen::MatrixXd& weights);
    double  evaluate (const Eigen::RowVectorXd& oneData, int classIdx) const;
    double  evaluate (const Eigen::MatrixXd& data, int dataIdx, int classIdx) const;
    Eigen::RowVectorXi     predict(const Eigen::RowVectorXd& oneData) const;

    double  evalError(const Eigen::MatrixXd& data, const Eigen::MatrixXi& labels, const Eigen::MatrixXd& weights) const;
    int     getLabelResponse(const Eigen::RowVectorXi& oneLabel, int classIdx) const;

    void output(std::ostream& out);

    const MLDecisionStumpParams* getParameters() const;
    int getTheBestClassIdx();

    void buildStump(const MLDecisionStumpParams& params);

private:
    std::auto_ptr<MLDecisionStumpImpl> impl_;

    friend class MLDecisionStumpImpl;
};

NAMESPACE_MACHINE_LEARNING_END
GLOBAL_NAMESPACE_END

#endif//ML_DECISION_STUMP_H_